import os
import sys
# 将项目根目录添加到模块搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import re
import pandas as pd
import pymysql
from datetime import datetime
from plate_business_filter import plate_business_filter_proc
from db_config import DB_CONFIG

data_collection_dir = 'd:/VERSION/stock_bankuai/data_collection/data4import'

# 连接数据库
def get_db_connection():
    return pymysql.connect(**DB_CONFIG)

# 创建所有表
def create_tables(conn):
    with conn.cursor() as cursor:
        # 创建已导入文件记录表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS imported_files (
            id INT AUTO_INCREMENT PRIMARY KEY,
            file_name VARCHAR(255) UNIQUE,
            import_time DATETIME
        )
        """)

        # 创建 A 股板块交易情况表（添加筛选、打分、推荐和验证相关字段）
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS `a_share_trading` (
            `id` INT NOT NULL AUTO_INCREMENT,
            `trade_date` VARCHAR(20) NULL DEFAULT NULL COLLATE 'utf8mb4_general_ci',
            `seq` INT NULL DEFAULT NULL COMMENT '序',
            `name` VARCHAR(255) NULL DEFAULT NULL COMMENT '名称' COLLATE 'utf8mb4_general_ci',
            `daily_gain` VARCHAR(20) NULL DEFAULT NULL COMMENT '涨幅%' COLLATE 'utf8mb4_general_ci',
            `three_day_gain` VARCHAR(20) NULL DEFAULT NULL COMMENT '3日涨幅%' COLLATE 'utf8mb4_general_ci',
            `gain_speed` VARCHAR(20) NULL DEFAULT NULL COMMENT '涨速%' COLLATE 'utf8mb4_general_ci',
            `leading_stock` VARCHAR(255) NULL DEFAULT NULL COMMENT '领涨股' COLLATE 'utf8mb4_general_ci',
            `up_down_count` VARCHAR(20) NULL DEFAULT NULL COMMENT '涨跌家数' COLLATE 'utf8mb4_general_ci',
            `up_down_ratio` VARCHAR(20) NULL DEFAULT NULL COMMENT '涨跌比' COLLATE 'utf8mb4_general_ci',
            `limit_up_count` INT NULL DEFAULT NULL COMMENT '涨停家数',
            `turnover_rate` VARCHAR(20) NULL DEFAULT NULL COMMENT '换手%' COLLATE 'utf8mb4_general_ci',
            `volume_ratio` VARCHAR(20) NULL DEFAULT NULL COMMENT '量比' COLLATE 'utf8mb4_general_ci',
            `three_day_turnover` VARCHAR(20) NULL DEFAULT NULL COMMENT '3日换手%' COLLATE 'utf8mb4_general_ci',
            `five_day_gain` VARCHAR(20) NULL DEFAULT NULL COMMENT '5日涨幅%' COLLATE 'utf8mb4_general_ci',
            `ten_day_gain` VARCHAR(20) NULL DEFAULT NULL COMMENT '10日涨幅%' COLLATE 'utf8mb4_general_ci',
            `twenty_day_gain` VARCHAR(20) NULL DEFAULT NULL COMMENT '20日涨幅%' COLLATE 'utf8mb4_general_ci',
            `trading_volume` VARCHAR(20) NULL DEFAULT NULL COMMENT '成交量' COLLATE 'utf8mb4_general_ci',
            `amount` VARCHAR(20) NULL DEFAULT NULL COMMENT '金额' COLLATE 'utf8mb4_general_ci',
            `total_market_cap` VARCHAR(20) NULL DEFAULT NULL COMMENT '总市值' COLLATE 'utf8mb4_general_ci',
            `circulating_market_cap` VARCHAR(20) NULL DEFAULT NULL COMMENT '流通市值' COLLATE 'utf8mb4_general_ci',
            `average_return` VARCHAR(20) NULL DEFAULT NULL COMMENT '平均收益' COLLATE 'utf8mb4_general_ci',
            `average_equity` VARCHAR(20) NULL DEFAULT NULL COMMENT '平均股本' COLLATE 'utf8mb4_general_ci',
            `pe_ratio` VARCHAR(20) NULL DEFAULT NULL COMMENT '市盈率' COLLATE 'utf8mb4_general_ci',
            `is_low_performance_eliminate` ENUM('是','否') NULL DEFAULT NULL COMMENT '是否为低绩效淘汰股' COLLATE 'utf8mb4_general_ci',
            `is_reserved_in_pool` ENUM('是','否') NULL DEFAULT NULL COMMENT '是否保留至候选池' COLLATE 'utf8mb4_general_ci',
            `fund_score_base` DECIMAL(5,1) NULL DEFAULT NULL COMMENT '资金维度基础分（0-10分）',
            `event_score_base` DECIMAL(5,1) NULL DEFAULT NULL COMMENT '事件维度基础分（0-10分）',
            `valuation_score_base` DECIMAL(5,1) NULL DEFAULT NULL COMMENT '估值维度基础分（0-10分）',
            `total_score` DECIMAL(5,1) NULL DEFAULT NULL COMMENT '总分',
            `is_recommended` ENUM('是','否') NULL DEFAULT NULL COMMENT '是否推荐' COLLATE 'utf8mb4_general_ci',
            `hit_status` VARCHAR(20) NULL DEFAULT NULL COMMENT '命中状态' COLLATE 'utf8mb4_general_ci',
	        `event_infos` MEDIUMTEXT NULL DEFAULT NULL COMMENT '事件的标题' COLLATE 'utf8mb4_general_ci',            PRIMARY KEY (`id`) USING BTREE,
            INDEX `trade_date` (`trade_date`) USING BTREE
        )
        COLLATE='utf8mb4_general_ci'
        ENGINE=InnoDB
        AUTO_INCREMENT=173
        ;

        """
        )

        # 创建时间对象表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS time_object (
            id INT AUTO_INCREMENT PRIMARY KEY,
            trading_date VARCHAR(20) NOT NULL COMMENT '交易日（YYYY-MM-DD）',
            data_status ENUM('待生成', '生成中', '已生成') COMMENT '数据状态',
            market_type VARCHAR(20)  COMMENT '(震荡市, 牛市, 熊市)',
            `hit_rate` FLOAT NULL DEFAULT NULL COMMENT '第二天板块涨幅的TOP20的命中率',
            UNIQUE KEY unique_trading_date (trading_date)
        )
        """)

        cursor.execute("""
        CREATE TABLE IF NOT EXISTS a_share_fund_flow (  
            id INT AUTO_INCREMENT PRIMARY KEY,  
            flow_date VARCHAR(20),  
            seq INT COMMENT '序',  
            code VARCHAR(20) COMMENT '代码',  
            name VARCHAR(255) COMMENT '名称',  
            latest_price VARCHAR(20) COMMENT '最新',  
            daily_gain VARCHAR(20) COMMENT '涨幅%',  
            main_net_inflow VARCHAR(20) COMMENT '主力净流入',  
            call_auction VARCHAR(20) COMMENT '集合竞价',  
            super_large_buy VARCHAR(20) COMMENT '超大单流入',  
            super_large_sell VARCHAR(20) COMMENT '超大单流出',  
            super_large_net VARCHAR(20) COMMENT '超大单净额',  
            super_large_net_ratio VARCHAR(20) COMMENT '超大单净占比%',  
            large_buy VARCHAR(20) COMMENT '大单流入',  
            large_sell VARCHAR(20) COMMENT '大单流出',  
            large_net VARCHAR(20) COMMENT '大单净额',  
            large_net_ratio VARCHAR(20) COMMENT '大单净占比%',  
            medium_buy VARCHAR(20) COMMENT '中单流入',  
            medium_sell VARCHAR(20) COMMENT '中单流出',  
            medium_net VARCHAR(20) COMMENT '中单净额',  
            medium_net_ratio VARCHAR(20) COMMENT '中单净占比%',  
            small_buy VARCHAR(20) COMMENT '小单流入',  
            small_sell VARCHAR(20) COMMENT '小单流出',  
            small_net VARCHAR(20) COMMENT '小单净额',  
            small_net_ratio VARCHAR(20) COMMENT '小单净占比%'  
        )
        """)

        cursor.execute("""
        CREATE TABLE IF NOT EXISTS policy_events (  
            id INT AUTO_INCREMENT PRIMARY KEY,  
            event_date VARCHAR(20),  
            time VARCHAR(255) COMMENT '时间',  
            title VARCHAR(255) COMMENT '标题',  
            hot_level VARCHAR(255) COMMENT '热门等级',  
            name VARCHAR(255) COMMENT '名称',  
            latest_price VARCHAR(255) COMMENT '最新',  
            daily_gain VARCHAR(255) COMMENT '涨幅%',  
            gain_speed VARCHAR(255) COMMENT '涨速%',  
            turnover_rate VARCHAR(255) COMMENT '换手%',  
            current_volume VARCHAR(255) COMMENT '现手',  
            total_volume VARCHAR(255) COMMENT '总手',  
            amount VARCHAR(255) COMMENT '金额',  
            total_market_cap VARCHAR(255) COMMENT '总市值',  
            circulating_market_cap VARCHAR(255) COMMENT '流通市值',  
            industry VARCHAR(255) COMMENT '所属行业',  
            three_day_gain VARCHAR(255) COMMENT '3日涨幅%',  
            six_day_gain VARCHAR(255) COMMENT '6日涨幅%',
            UNIQUE KEY unique_title (title)
        )
        """)

        cursor.execute("""
        CREATE TABLE IF NOT EXISTS hs300_index (  
            `id` INT NOT NULL AUTO_INCREMENT,
            `index_date` VARCHAR(20) NULL DEFAULT NULL COMMENT '日期' COLLATE 'utf8mb4_general_ci',
            `closing_index` VARCHAR(20) NULL DEFAULT NULL COMMENT '收盘价格' COLLATE 'utf8mb4_general_ci',
            `fund_weight` INT NULL DEFAULT NULL COMMENT '资金权重',
            `event_weight` INT NULL DEFAULT NULL COMMENT '事件权重',
            `valuation_weight` INT NULL DEFAULT NULL COMMENT '估值权重',
            `market_trend` VARCHAR(20) NULL DEFAULT NULL COMMENT '市场趋势：牛市、熊市、震荡市' COLLATE 'utf8mb4_general_ci',
            PRIMARY KEY (`id`) USING BTREE,
            INDEX `index_date` (`index_date`) USING BTREE
        )
        """)
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS industry_hedge_relations (
                id INT PRIMARY KEY AUTO_INCREMENT,
                plate VARCHAR(100) NOT NULL,
                hedge_plate VARCHAR(100) NOT NULL,
                correlation_coefficient DECIMAL(5,3) NOT NULL,
                logic TEXT,
                INDEX idx_plate (plate),  -- 优化查询效率
                INDEX idx_hedge_plate (hedge_plate)
        )
        """)

    conn.commit()

# 检查文件是否已导入
def is_file_imported(conn, file_name):
    with conn.cursor() as cursor:
        cursor.execute("SELECT 1 FROM imported_files WHERE file_name = %s", (file_name,))
        return cursor.fetchone() is not None

# 记录文件已导入
def mark_file_imported(conn, file_name):
    with conn.cursor() as cursor:
        cursor.execute("INSERT INTO imported_files (file_name, import_time) VALUES (%s, %s)", 
                      (file_name, datetime.now()))
    conn.commit()

# 从文件名中提取日期
def extract_date_from_filename(file_name):
    # 修改正则表达式，去掉多余的反斜杠
    match = re.search(r'\d{4}-\d{2}-\d{2}', file_name)
    if match:
        return match.group()
    return None

# 导入 A 股板块交易情况数据
# 修改导入函数，添加异常处理，以 import_a_share_trading 为例
def import_a_share_trading(conn, file_path):
    file_name = os.path.basename(file_path)
    if is_file_imported(conn, file_name):
        return
    try:
        df = pd.read_excel(file_path, engine='openpyxl')
        trade_date = extract_date_from_filename(file_name)
        with conn.cursor() as cursor:
            for _, row in df.iterrows():
                cursor.execute("""
                INSERT INTO a_share_trading (  
                    trade_date, seq, name, daily_gain, three_day_gain, gain_speed,  
                    leading_stock, up_down_count, up_down_ratio, limit_up_count,  
                    turnover_rate, volume_ratio, three_day_turnover, five_day_gain,  
                    ten_day_gain, twenty_day_gain, trading_volume, amount,  
                    total_market_cap, circulating_market_cap, average_return,  
                    average_equity, pe_ratio
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """, (
                    trade_date, row['序'], row['名称'].strip(), str(row['涨幅%']), str(row['3日涨幅%']), str(row['涨速%']),  
                    row['领涨股'], str(row['涨跌家数']), str(row['涨跌比']), row['涨停家数'],  
                    str(row['换手%']), str(row['量比']), str(row['3日换手%']), str(row['5日涨幅%']),  
                    str(row['10日涨幅%']), str(row['20日涨幅%']), str(row['成交量']), str(row['金额']),  
                    str(row['总市值']), str(row['流通市值']), str(row['平均收益']),  
                    str(row['平均股本']), str(row['市盈率'])
                ))
        conn.commit()
        mark_file_imported(conn, file_name)
        
        # 给时间对象表插入数据
        with conn.cursor() as cursor:
            cursor.execute(
                "INSERT INTO time_object (trading_date, data_status, market_type) VALUES (%s, %s, %s)"
                " ON DUPLICATE KEY UPDATE data_status = VALUES(data_status), market_type = VALUES(market_type)",
                (trade_date, '已生成', None)
            )
        conn.commit()
    except Exception as e:
        print(f"导入文件 {file_path} 时出错: {str(e)}")
        conn.rollback()

# 导入 A 股板块资金流向数据
def import_a_share_fund_flow(conn, file_path):
    file_name = os.path.basename(file_path)
    if is_file_imported(conn, file_name):
        return
    try:
        df = pd.read_excel(file_path, engine='openpyxl')
        flow_date = extract_date_from_filename(file_name)
        if not flow_date:
            print(f"错误：无法从文件名{file_name}中提取日期")
            return
        with conn.cursor() as cursor:
            for _, row in df.iterrows():
                cursor.execute("""
                INSERT INTO a_share_fund_flow (
                    flow_date, seq, code, name, latest_price, daily_gain, main_net_inflow, 
                    call_auction, super_large_buy, super_large_sell, super_large_net, 
                    super_large_net_ratio, large_buy, large_sell, large_net, 
                    large_net_ratio, medium_buy, medium_sell, medium_net, 
                    medium_net_ratio, small_buy, small_sell, small_net, 
                    small_net_ratio
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """, (
                    flow_date, row['序'], row['代码'], row['名称'].strip(), str(row['最新']), str(row['涨幅%']), str(row['主力净流入']), 
                    str(row['集合竞价']), str(row['超大单流入']), str(row['超大单流出']), str(row['超大单净额']), 
                    str(row['超大单净占比%']), str(row['大单流入']), str(row['大单流出']), str(row['大单净额']), 
                    str(row['大单净占比%']), str(row['中单流入']), str(row['中单流出']), str(row['中单净额']), 
                    str(row['中单净占比%']), str(row['小单流入']), str(row['小单流出']), str(row['小单净额']), 
                    str(row['小单净占比%'])
                ))
        conn.commit()
        mark_file_imported(conn, file_name)
        
        # 给时间对象表插入数据
        with conn.cursor() as cursor:
            cursor.execute(
                "INSERT INTO time_object (trading_date, data_status, market_type) VALUES (%s, %s, %s)"
                " ON DUPLICATE KEY UPDATE data_status = VALUES(data_status), market_type = VALUES(market_type)",
                (flow_date, '已生成', None)
            )
        conn.commit()
    except Exception as e:
        print(f"导入资金流向文件 {file_path} 时出错: {str(e)}")
        conn.rollback()

# 导入政策和事件数据
def import_policy_events(conn, file_path):
    file_name = os.path.basename(file_path)
    if is_file_imported(conn, file_name):
        return
    
    df = pd.read_excel(file_path, engine='openpyxl')
    event_date = extract_date_from_filename(file_name)
    
    with conn.cursor() as cursor:
        for _, row in df.iterrows():

            # 从 row['时间'] 提取日期并转换为 yyyy-mm-dd 格式
            try:
                current_year = datetime.now().year
                time_str = f'{current_year}-{row["时间"].strip()}'
                event_date1 = datetime.strptime(time_str, '%Y-%m-%d %H:%M').strftime('%Y-%m-%d')
            except ValueError:
                print(f'日期格式错误: {row["时间"]}')
                continue
            cursor.execute(
                """
                INSERT IGNORE INTO policy_events (
                    event_date, time, title, hot_level, name, latest_price, 
                    daily_gain, gain_speed, turnover_rate, current_volume, 
                    total_volume, amount, total_market_cap, circulating_market_cap, 
                    industry, three_day_gain, six_day_gain
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                """,
                (
                    event_date1, row['时间'], row['标题'], row['热门等级'], row['名称'].strip(), str(row['最新']), 
                    str(row['涨幅%']), str(row['涨速%']), str(row['换手%']), str(row['现手']), 
                    str(row['总手']), str(row['金额']), str(row['总市值']), str(row['流通市值']), 
                    row['所属行业'].strip(), str(row['3日涨幅%']), str(row['6日涨幅%'])
                )
            )
    conn.commit()
    mark_file_imported(conn, file_name)
    
    conn.commit()

# 导入沪深 300 指数数据
def import_hs300_index(conn, file_path):
    try:
        # 添加文件验证
        if not os.path.exists(file_path):
            print(f"错误：文件{file_path}不存在")
            return
        if os.path.getsize(file_path) == 0:
            print(f"错误：文件{file_path}为空")
            return
        df = pd.read_excel(file_path, engine='openpyxl')
    except Exception as e:
        print(f"导入HS300指数文件 {file_path} 时出错: {str(e)}")
        conn.rollback()
        return

    # 检查必要列是否存在
    required_columns = ['日期', '收盘指数']
    for col in required_columns:
        if col not in df.columns:
            print(f"错误：文件 {file_path} 中缺少必要列 {col}")
            return
    
    # 先获取所有日期，一次性删除旧数据，提升性能
    dates = df['日期'].astype(str).tolist()
    if dates:
        with conn.cursor() as cursor:
            # 删除旧数据
            delete_query = f"DELETE FROM hs300_index"
            cursor.execute(delete_query)
            
            # 插入新数据
            for _, row in df.iterrows():
                index_date = row['日期']
                closing_index = str(row['收盘指数'])
                cursor.execute("INSERT INTO hs300_index (index_date, closing_index) VALUES (%s, %s)", 
                              (index_date, closing_index))
    conn.commit()
    # 移除记录文件已导入的步骤

# 主函数
def import_new_data(conn):

    try:
        create_tables(conn)
        
        # 遍历目录下的文件
        for root, _, files in os.walk(data_collection_dir):
            for file in files:
                file_path = os.path.join(root, file)
                if file.startswith('bankuai-') and not file.startswith('bankuai-zijin-') and file.endswith(('.xlsx', '.XLSX')):
                    import_a_share_trading(conn, file_path)
                elif file.startswith('bankuai-zijin-') and file.endswith(('.xlsx', '.XLSX')):
                    import_a_share_fund_flow(conn, file_path)
                elif file.startswith('shijian-') and file.endswith(('.xlsx', '.XLSX')):
                    import_policy_events(conn, file_path)
                elif file == 'HS300.XLSX' or file == 'HS300.xlsx':
                    import_hs300_index(conn, file_path)
    finally:
        conn.close()
        print(f"导入完成！")

if __name__ == '__main__':

    import_new_data(get_db_connection())
    conn = get_db_connection()
    # 调用板块计算和推荐
    try:
        
        while True:
            date_input = input('请输入日期（格式：yyyy-mm-dd，直接回车使用最新日期）: ')
            if not date_input:
                plate_business_filter_proc(conn)
                break
            if re.match(r'^\d{4}-\d{2}-\d{2}$', date_input):
                plate_business_filter_proc(conn, last_date=date_input)
                break
            else:
                print('日期格式不正确，请使用 yyyy-mm-dd 格式重新输入。')
        conn.close()
    except Exception as e:
        print(f"调用 plate_business_filter 时出错: {str(e)}")